import networkx as nx
import matplotlib.pyplot as plt
import pulp
import numpy as np
from matplotlib.lines import Line2D
import random
import math
from collections import deque
from matplotlib.font_manager import FontProperties
font = FontProperties(fname='resources/SimHei.ttf', size=12)  # 替换为实际字体路径

# 设置随机种子保证结果可复现
np.random.seed(42)
random.seed(42)

# ======================
# 1. 管网拓扑生成（确保正确流向）
# ======================
def generate_gas_network(num_levels=4, nodes_per_level=5, num_sources=2, num_leaks=3, num_customers=15):
    """生成燃气管网拓扑，确保正确的流向（气源->阀门->用户）"""
    G = nx.DiGraph()
    positions = {}
    
    # 创建分层结构确保流向正确
    levels = {}
    node_id = 0
    
    # 第0层：气源
    sources = []
    for i in range(num_sources):
        node_id_str = f"S{i}"
        G.add_node(node_id_str, type='source', pressure=0.8 + random.uniform(0.1, 0.3))
        # 气源放在左侧
        positions[node_id_str] = (0, random.uniform(0.2, 0.8))
        sources.append(node_id_str)
        levels[node_id_str] = 0
        node_id += 1
    
    # 中间层：阀门和交汇点
    for level in range(1, num_levels - 1):
        level_nodes = []
        for j in range(nodes_per_level):
            # 随机决定是阀门还是交汇点
            if random.random() < 0.4:  # 40%的概率是阀门
                node_id_str = f"V{node_id}"
                G.add_node(node_id_str, type='valve', status='open', 
                          operation_time=random.uniform(1, 5))
                level_nodes.append(node_id_str)
            else:  # 60%的概率是交汇点
                node_id_str = f"J{node_id}"
                G.add_node(node_id_str, type='junction', pressure=0.3 + random.uniform(0.1, 0.4))
                level_nodes.append(node_id_str)
            
            # 位置：x坐标按层级递增
            positions[node_id_str] = (level/(num_levels-1), random.uniform(0, 1))
            levels[node_id_str] = level
            node_id += 1
    
    # 最后一层：用户和泄漏点
    customers = []
    leaks = []
    for j in range(num_customers + num_leaks):
        # 用户和泄漏点放在右侧
        x_pos = 1.0
        
        # 80%概率是用户，20%是泄漏点
        if j < num_customers:
            node_id_str = f"C{node_id}"
            G.add_node(node_id_str, type='customer', pressure=0.2 + random.uniform(0.05, 0.15),
                      num_users=random.randint(10, 100))
            customers.append(node_id_str)
        else:
            node_id_str = f"L{node_id}"
            G.add_node(node_id_str, type='leak', pressure=0.2 + random.uniform(0.1, 0.3),
                      leak_rate=10 + random.uniform(0, 20))
            leaks.append(node_id_str)
        
        positions[node_id_str] = (x_pos, random.uniform(0, 1))
        levels[node_id_str] = num_levels - 1
        node_id += 1
    
    # 添加边（确保流向从气源到用户）
    all_nodes = list(G.nodes())
    
    # 连接气源到第一层
    for source in sources:
        for node in all_nodes:
            if levels.get(node, -1) == 1:
                # 只连接在右侧的节点
                if positions[source][0] < positions[node][0]:
                    length = np.linalg.norm(np.array(positions[source]) - np.array(positions[node]))
                    diameter = random.choice([50, 100, 150, 200])
                    G.add_edge(source, node, length=length, diameter=diameter, 
                              type='pipeline', flow_direction='forward')
    
    # 连接中间层
    for level in range(1, num_levels - 1):
        current_level_nodes = [n for n in all_nodes if levels.get(n, -1) == level]
        next_level_nodes = [n for n in all_nodes if levels.get(n, -1) == level + 1]
        
        for u in current_level_nodes:
            for v in next_level_nodes:
                # 确保流向从左到右
                if positions[u][0] < positions[v][0]:
                    # 随机连接（避免全连接）
                    if random.random() < 0.6:
                        length = np.linalg.norm(np.array(positions[u]) - np.array(positions[v]))
                        diameter = random.choice([50, 100, 150, 200])
                        G.add_edge(u, v, length=length, diameter=diameter, 
                                  type='pipeline', flow_direction='forward')
    
    # 连接最后一层
    last_level_nodes = [n for n in all_nodes if levels.get(n, -1) == num_levels - 2]
    for u in last_level_nodes:
        for v in customers + leaks:
            # 确保流向从左到右
            if positions[u][0] < positions[v][0]:
                # 随机连接
                if random.random() < 0.7:
                    length = np.linalg.norm(np.array(positions[u]) - np.array(positions[v]))
                    diameter = random.choice([50, 100, 150, 200])
                    G.add_edge(u, v, length=length, diameter=diameter, 
                              type='pipeline', flow_direction='forward')
    
    # 添加一些横向连接增加网络复杂性
    for level in range(1, num_levels - 1):
        level_nodes = [n for n in all_nodes if levels.get(n, -1) == level]
        for i in range(len(level_nodes)):
            for j in range(i + 1, len(level_nodes)):
                if random.random() < 0.3:
                    u, v = level_nodes[i], level_nodes[j]
                    # 随机决定方向，但保持从左到右
                    if positions[u][0] < positions[v][0]:
                        from_node, to_node = u, v
                    else:
                        from_node, to_node = v, u
                    
                    length = np.linalg.norm(np.array(positions[from_node]) - np.array(positions[to_node]))
                    diameter = random.choice([50, 100, 150, 200])
                    G.add_edge(from_node, to_node, length=length, diameter=diameter, 
                              type='pipeline', flow_direction='forward')
    
    # 收集所有阀门节点
    valves = [n for n in G.nodes() if G.nodes[n].get('type') == 'valve']
    
    return G, positions, sources, leaks, valves, customers

# 生成管网
G, positions, sources, leaks, valves, customers = generate_gas_network()

# ======================
# 2. 最小割建模 (隔离泄漏点)
# ======================
def build_min_cut_model(G, sources, leaks, positions):
    """构建最小割模型隔离泄漏点"""
    # 创建副本图
    H = nx.DiGraph()
    
    # 添加所有节点
    for node in G.nodes():
        H.add_node(node, **G.nodes[node])
    
    # 添加所有边
    for u, v, data in G.edges(data=True):
        H.add_edge(u, v, **data)
    
    # 添加超级源点和超级汇点
    super_source = "SUPER_SOURCE"
    super_sink = "SUPER_SINK"
    H.add_node(super_source)
    H.add_node(super_sink)
    
    # 连接超级源点到所有气源
    for source in sources:
        H.add_edge(super_source, source, capacity=float('inf'))
    
    # 连接所有泄漏点到超级汇点
    for leak in leaks:
        H.add_edge(leak, super_sink, capacity=float('inf'))
    
    # 设置阀门边的容量（关闭代价）
    for node in valves:
        if H.nodes[node].get('status') == 'open':
            # 获取阀门操作时间作为基础容量
            operation_time = H.nodes[node].get('operation_time', 3.0)
            
            # 计算下游影响 - 通过BFS找到所有下游节点
            downstream_nodes = set()
            queue = deque([node])
            while queue:
                current = queue.popleft()
                for successor in H.successors(current):
                    if successor not in downstream_nodes:
                        downstream_nodes.add(successor)
                        queue.append(successor)
            
            # 计算受影响用户数
            affected_users = 0
            for n in downstream_nodes:
                if H.nodes[n].get('type') == 'customer':
                    affected_users += H.nodes[n].get('num_users', 0)
            
            # 阀门容量 = 操作时间 * (1 + 受影响用户/100)
            capacity = operation_time * (1 + affected_users / 100)
            
            # 在图中标记阀门容量
            H.nodes[node]['capacity'] = capacity
    
    # 转换阀门节点为边（用于最小割计算）
    for valve in valves:
        if valve in H.nodes():
            # 获取所有前驱和后继
            predecessors = list(H.predecessors(valve))
            successors = list(H.successors(valve))
            
            # 移除阀门节点
            H.remove_node(valve)
            
            # 添加新边连接前驱和后继，容量为阀门容量
            for u in predecessors:
                for v in successors:
                    capacity = H.nodes.get(valve, {}).get('capacity', 10.0)
                    H.add_edge(u, v, capacity=capacity, valve_edge=True, original_valve=valve)
    
    # 设置其他边的容量为无限大（不能切割）
    for u, v in list(H.edges()):
        if 'capacity' not in H.edges[(u, v)]:
            H.edges[(u, v)]['capacity'] = float('inf')
    
    # 计算最小割
    try:
        cut_value, partition = nx.minimum_cut(H, super_source, super_sink)
        reachable, non_reachable = partition
        
        # 收集需要关闭的阀门
        valves_to_close = set()
        for u, v, data in H.edges(data=True):
            if data.get('valve_edge', False):
                if (u in reachable and v in non_reachable):
                    valve_name = data.get('original_valve', 'unknown')
                    valves_to_close.add(valve_name)
        
        return list(valves_to_close), cut_value
    except nx.NetworkXUnbounded:
        print("警告：最小割计算失败，使用备用方法")
        # 备用方法：关闭所有泄漏点上游的阀门
        valves_to_close = set()
        for leak in leaks:
            # 找到所有从气源到泄漏点的路径
            for source in sources:
                try:
                    paths = nx.all_simple_paths(G, source, leak)
                    for path in paths:
                        # 找到路径上的阀门
                        for node in path:
                            if node in valves:
                                valves_to_close.add(node)
                except nx.NetworkXNoPath:
                    continue
        
        return list(valves_to_close), len(valves_to_close) * 10.0

# 计算最小割
min_cut_valves, min_cut_value = build_min_cut_model(G, sources, leaks, positions)
print(f"最小割算法建议关闭阀门: {min_cut_valves}")
print(f"最小割值(代价估计): {min_cut_value:.2f}")

# ======================
# 3. 多目标优化 (线性规划)
# ======================
def optimize_valve_selection(G, valves, leaks, customers, min_cut_valves, positions):
    """使用线性规划优化阀门选择"""
    # 创建问题实例
    prob = pulp.LpProblem("Valve_Optimization", pulp.LpMinimize)
    
    # 决策变量 (是否关闭阀门)
    valve_vars = pulp.LpVariable.dicts("Valve", valves, cat='Binary')
    
    # 目标函数权重
    w_leak = 0.5  # 泄漏总量权重
    w_users = 0.3  # 用户影响权重
    w_time = 0.2   # 操作时间权重
    
    # 1. 泄漏总量目标
    leak_objective = 0
    for leak in leaks:
        # 获取泄漏率
        leak_rate = G.nodes[leak].get('leak_rate', 15)
        
        # 计算该泄漏点的隔离时间（取决于阀门操作时间）
        isolation_time = pulp.LpAffineExpression()
        for valve in min_cut_valves:
            if valve in valves:
                # 计算阀门到泄漏点的距离（作为时间因子）
                dist = np.linalg.norm(np.array(positions.get(valve, (0,0))) - 
                                     np.array(positions.get(leak, (0,0))))
                time_factor = max(0.1, dist)  # 避免除零
                
                # 阀门关闭对泄漏的影响
                isolation_time += valve_vars[valve] * G.nodes[valve].get('operation_time', 3) * time_factor
        
        # 泄漏总量 = 泄漏率 × 隔离时间
        leak_objective += leak_rate * isolation_time
    
    # 2. 用户影响目标
    user_objective = 0
    for customer in customers:
        # 检查客户是否会被影响（是否在关闭阀门的下游）
        is_affected = pulp.LpAffineExpression()
        for valve in min_cut_valves:
            if valve in valves:
                # 检查从阀门到客户是否有路径
                try:
                    # 简化的下游检查：x坐标在阀门右侧
                    if positions.get(customer, (0,0))[0] > positions.get(valve, (0,0))[0]:
                        is_affected += valve_vars[valve]
                except:
                    pass
        
        # 用户影响 = 用户数 × 是否受影响
        user_objective += G.nodes[customer].get('num_users', 30) * is_affected
    
    # 3. 操作时间目标 (最大操作时间)
    max_time = pulp.LpVariable("Max_Time", lowBound=0)
    for valve in min_cut_valves:
        if valve in valves:
            # 阀门操作时间
            op_time = G.nodes[valve].get('operation_time', 3)
            
            # 考虑阀门位置（距离气源的距离）
            dist_to_source = min([np.linalg.norm(np.array(positions.get(valve, (0,0))) - 
                                     np.array(positions.get(source, (0,0)))) 
                             for source in sources])
            time_factor = max(0.1, dist_to_source)  # 避免除零
            
            prob += max_time >= valve_vars[valve] * op_time * time_factor
    
    # 总目标函数
    prob += w_leak * leak_objective + w_users * user_objective + w_time * max_time
    
    # 约束：必须隔离所有泄漏点
    for leak in leaks:
        # 至少关闭一个能隔离该泄漏点的阀门
        prob += pulp.lpSum(valve_vars[valve] for valve in min_cut_valves if valve in valves) >= 1
    
    # 求解问题
    prob.solve(pulp.PULP_CBC_CMD(msg=True))
    
    # 收集结果
    selected_valves = [v for v in valves if valve_vars[v].value() == 1]
    
    # 计算目标值
    leak_total = pulp.value(leak_objective)
    user_affected = pulp.value(user_objective)
    max_operation_time = pulp.value(max_time)
    
    return selected_valves, leak_total, user_affected, max_operation_time

# 运行优化
optimized_valves, leak_total, user_affected, max_time = optimize_valve_selection(
    G, valves, leaks, customers, min_cut_valves, positions
)
print(f"优化后关闭阀门: {optimized_valves}")
print(f"预计泄漏总量: {leak_total:.2f} 单位")
print(f"受影响的用户数: {user_affected:.0f}")
print(f"最大操作时间: {max_time:.2f} 分钟")

# ======================
# 4. 可视化
# ======================
def visualize_network(G, positions, sources, leaks, valves, customers, optimized_valves):
    """可视化管网和优化结果"""
    plt.figure(figsize=(16, 10))
    
    # 创建节点颜色和大小映射
    node_colors = []
    node_sizes = []
    labels = {}
    
    # 绘制节点
    for node in G.nodes():
        node_type = G.nodes[node].get('type', 'junction')
        
        if node in sources:
            node_colors.append('green')
            node_sizes.append(300)
            labels[node] = 'S'  # 气源
        elif node in leaks:
            node_colors.append('red')
            node_sizes.append(300)
            labels[node] = 'L'  # 泄漏点
        elif node in valves:
            if node in optimized_valves:
                node_colors.append('purple')
                node_sizes.append(300)
            else:
                node_colors.append('blue')
                node_sizes.append(300)
            labels[node] = 'V'  # 阀门
        elif node in customers:
            node_colors.append('orange')
            node_sizes.append(300)
            labels[node] = 'C'  # 用户
        else:
            node_colors.append('gray')
            node_sizes.append(300)
            labels[node] = 'O'  # 其他节点
    
    # 绘制节点
    nx.draw_networkx_nodes(G, positions, node_color=node_colors, node_size=node_sizes)
    
    # 绘制边（带箭头表示流向）
    for u, v in G.edges():
        # 跳过连接到超级节点的边
        if u.startswith("SUPER_") or v.startswith("SUPER_"):
            continue
        
        # 绘制箭头
        dx = positions[v][0] - positions[u][0]
        dy = positions[v][1] - positions[u][1]
        distance = math.sqrt(dx**2 + dy**2)
        
        # 归一化
        if distance > 0:
            dx, dy = dx/distance, dy/distance
        
        # 缩短箭头避免重叠
        arrow_length = distance * 0.8
        arrow_start = (positions[u][0] + dx * distance * 0.15, 
                      positions[u][1] + dy * distance * 0.15)
        arrow_end = (positions[u][0] + dx * arrow_length, 
                    positions[u][1] + dy * arrow_length)
        
        # 绘制线
        plt.plot([arrow_start[0], arrow_end[0]], 
                 [arrow_start[1], arrow_end[1]], 
                 'k-', alpha=0.7, linewidth=1.5)
        
        # 绘制箭头头
        arrow_head_length = max(0.02, distance * 0.1)
        plt.arrow(arrow_end[0], arrow_end[1], 
                 dx * arrow_head_length, dy * arrow_head_length, 
                 head_width=0.006, head_length=0.006, 
                 fc='k', ec='k', alpha=0.7)
    
    # 添加标签
    for node, pos in positions.items():
        if labels.get(node, '') != '':
            plt.text(pos[0], pos[1] + 0.03, labels[node], 
                    fontsize=9, ha='center', va='center',
                    bbox=dict(facecolor='white', alpha=0.7, boxstyle='round,pad=0.2'))
    
    # 添加图例
    legend_elements = [
        Line2D([0], [0], marker='o', color='w', markerfacecolor='green', markersize=10, label='气源 (S)'),
        Line2D([0], [0], marker='o', color='w', markerfacecolor='red', markersize=10, label='泄漏点 (L)'),
        Line2D([0], [0], marker='o', color='w', markerfacecolor='blue', markersize=10, label='阀门-开启 (V)'),
        Line2D([0], [0], marker='o', color='w', markerfacecolor='purple', markersize=10, label='阀门-关闭 (V)'),
        Line2D([0], [0], marker='o', color='w', markerfacecolor='orange', markersize=10, label='用户 (C)'),
        Line2D([0], [0], marker='o', color='w', markerfacecolor='gray', markersize=7, label='管网节点'),
        Line2D([0], [0], color='k', linestyle='-', lw=2, label='天然气流向'),
    ]
    
    plt.legend(handles=legend_elements, loc='upper left', fontsize=10, prop=font)
    
    # 添加标题
    plt.title("燃气管网泄漏控制优化", fontsize=16, pad=20, fontproperties=font)
    
    # 添加结果摘要
    result_text = (
        f"优化结果:\n"
        f"- 关闭阀门: {len(optimized_valves)}个 ({', '.join(optimized_valves[:3])}{'...' if len(optimized_valves) > 3 else ''})\n"
        f"- 受影响用户: {int(user_affected)}人\n"
        f"- 泄漏总量: {leak_total:.1f}单位\n"
        f"- 最大操作时间: {max_time:.1f}分钟\n"
        f"- 最小割值: {min_cut_value:.1f}"
    )
    
    plt.figtext(0.15, 0.15, result_text, 
               ha="left", va="bottom", fontsize=12,
               bbox=dict(facecolor='white', alpha=0.8, boxstyle='round,pad=0.5'), fontproperties=font)
    
    plt.axis('off')
    plt.tight_layout()
    plt.savefig('gas_network_optimization.png', dpi=300, bbox_inches='tight')
    plt.show()

# 可视化结果
visualize_network(G, positions, sources, leaks, valves, customers, optimized_valves)